#!/usr/bin/env python
import copy
import rosbag
import os
import os.path as osp
import yaml
import rospkg
import shutil
from sensor_msgs import point_cloud2
import cv2
import cv_bridge
import numpy as np
import pcl
import concurrent.futures


class BagDealer(object):
    def __init__(self, bag_path, config_path):
        self.bag_path = bag_path
        self.config_path = config_path
        self.bag_file = osp.join(self.bag_path, os.listdir(self.bag_path)[0])
        self.config_file = osp.join(self.config_path, os.listdir(self.config_path)[0])
        self.data_dir = osp.join(self.bag_path, 'data')

    @property
    def m_bag(self):
        return rosbag.Bag(self.bag_file, 'r')

    @property
    def m_config(self):
        with open(self.config_file, 'r') as f:
            return yaml.load(f)

    def extract(self, target_path = None, override = False):
        if override:
            if osp.exists(self.data_dir):
                print("refresh data directory %s" % target_path)
                shutil.rmtree(self.data_dir)
            os.mkdir(self.data_dir)
        else:
            if osp.exists(self.data_dir) == False:
                os.mkdir(self.data_dir)
        img_topics = []
        pc_topics = []
        ls_topics = []
        for k,v in self.m_config.items():
            if isinstance(v, list):
                img_topics.extend(v)
            elif isinstance(v, str) and v != "" and 'points' in v:
                pc_topics.append(v)
        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
            executor.map(self._extract_image, img_topics)
        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
            executor.map(self._extract_pointcloud2, pc_topics)

    def _extract_image(self, topic):
        assert self.data_dir, "data directory is None"
        target_dir = osp.join(self.data_dir, topic.split('/')[1])
        if osp.exists(target_dir) is False:
            os.mkdir(target_dir)
        else:
            shutil.rmtree(target_dir)
            os.mkdir(target_dir)
        bridge = cv_bridge.CvBridge()
        for topic, msg, t in self.m_bag.read_messages(topics=[topic]):
            img = bridge.imgmsg_to_cv2(msg)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            cv2.imwrite(osp.join(target_dir, str(t)+'.jpg'), img, [int(cv2.IMWRITE_JPEG_QUALITY),100])

    def _extract_pointcloud2(self, topic):
        assert self.data_dir, "data directory is None"
        target_dir = osp.join(self.data_dir, topic.split('/')[1])
        if osp.exists(target_dir) is False:
            os.mkdir(target_dir)
        else:
            shutil.rmtree(target_dir)
            os.mkdir(target_dir)
        for topic, msg, t in self.m_bag.read_messages(topics=[topic]):
            gen = point_cloud2.read_points(msg, field_names=('x', 'y', 'z', 'intensity'))
            data = np.zeros(0)
            flag = True
            for idx, p in enumerate(gen):
                if flag:
                    data = np.array([p[0], p[1], p[2], p[3]], dtype=np.float32).reshape(-1, 4)
                    flag = False
                else:
                    data = np.vstack((data, np.array([p[0], p[1], p[2], p[3]], dtype=np.float32)))
            pc = pcl.PointCloud_PointXYZI(data)
            pcl.save(pc, osp.join(target_dir, str(t)+'.pcd'))

    def align_sensors(self, tollerance = 0.05):
        topic_dirs = os.listdir(self.data_dir)
        topic_maps = {k: sorted(os.listdir(osp.join(self.data_dir, k)))
                      for k in topic_dirs if os.listdir(osp.join(self.data_dir, k)) != []}
        topic_list, file_list = [], []
        for topic, files in topic_maps.items():
            topic_list.append(topic)
            file_list.append(files)
        ts_list = copy.deepcopy(file_list)
        for i, lst in enumerate(file_list):
            ts_list[i] = [float(int(int(j.rsplit('.')[0]) / 1e6)) / 1e3 for j in lst]
        print ts_list[0]
        print ts_list[1]
        fptr1, fptr2 = 0, 0
        f1_len, f2_len = len(ts_list[0]), len(ts_list[1])
        done_flag = False
        while not done_flag:
            diff = ts_list[0][fptr1] - ts_list[1][fptr2]
            if abs(diff <= tollerance):
                fptr1 += 1
                fptr2 += 1
            elif diff < -tollerance:
                fptr1 += 1
                self.del_data_file(topic_list[1], file_list[1][fptr2])
                del ts_list[1][fptr2]
            elif diff > tollerance:
                fptr2 += 1
                self.del_data_file(topic_list[0], file_list[0][fptr1])
                del ts_list[0][fptr1]
            if fptr1 == len(ts_list[0]) - 1:
                done_flag = True
                for i in range(fptr2, len(ts_list[1])):
                    self.del_data_file(topic_list[1], file_list[1][fptr2])
            if fptr2 == len(ts_list[1]) - 1:
                done_flag = True
                for i in range(fptr1, len(ts_list[0])):
                    self.del_data_file(topic_list[0], file_list[0][fptr1])



    def del_data_file(self, topic, file_name):
        if osp.exists(osp.join(self.data_dir, topic, file_name)):
            os.remove(osp.join(self.data_dir, topic, file_name))


if __name__== "__main__":
    base_path = None
    for path in rospkg.get_ros_package_path().split(':'):
        if 'yunfei' in path:
            base_path = path
    assert base_path, "base path not set correctly"
    base_path = osp.join(base_path, 'bag_dealer')
    bag_path = osp.join(base_path, 'bags')
    config_path = osp.join(base_path, 'config')
    bd = BagDealer(bag_path, config_path)
    bd.extract()
    # bd.align_sensors()